/*
 * Copyright 2023-2024 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.ai.chat.prompt;

import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.util.I18nUtil;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.model.Media;
import org.springframework.core.io.Resource;

import java.io.IOException;
import java.nio.charset.Charset;
import java.util.*;
import java.util.Map.Entry;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class PromptTemplate implements PromptTemplateActions, PromptTemplateMessageActions {
    private static final String DEFAULT_START_DELIMITER = "{{";
    private static final String DEFAULT_END_DELIMITER = "}}";
    protected String template;

    private final String startDelimiter;
    private final String endDelimiter;

    private Map<String, Object> dynamicModel = new HashMap<>();

    public PromptTemplate(String template) {
        this(template, DEFAULT_START_DELIMITER, DEFAULT_END_DELIMITER);
    }

    public PromptTemplate(String template, String startDelimiter, String endDelimiter) {
        this.template = template;
        this.startDelimiter = startDelimiter;
        this.endDelimiter = endDelimiter;
    }

    public PromptTemplate(String template, Map<String, Object> model) {
        this(template);
        // If the template string is not valid, an exception will be thrown
        try {
            for (Entry<String, Object> entry : model.entrySet()) {
                add(entry.getKey(), entry.getValue());
            }
        }
        catch (Exception ex) {
            throw new IllegalArgumentException("The template string is not valid.", ex);
        }
    }

    public PromptTemplate(Resource resource, Map<String, Object> model) throws IOException {
        this.template = template;
        this.startDelimiter = DEFAULT_START_DELIMITER;
        this.endDelimiter = DEFAULT_END_DELIMITER;
        // If the template string is not valid, an exception will be thrown
        try {
            for (Entry<String, Object> entry : model.entrySet()) {
                add(entry.getKey(), entry.getValue());
            }
        }
        catch (Exception ex) {
            throw new IllegalArgumentException("The template string is not valid.", ex);
        }
    }

    public void add(String name, Object value) {
        this.dynamicModel.put(name, value);
    }

    public String getTemplate() {
        return this.template;
    }

    @Override
    public String render() {
        return render(this.dynamicModel);
    }

    @Override
    public String render(Map<String, Object> model) {
//        validate(model);
        String result = template;
        Locale locale = I18nUtil.getLocale();
        for (Entry<String, Object> entry : model.entrySet()) {
            String placeholder = startDelimiter + entry.getKey() + endDelimiter;
            String value = entry.getValue() instanceof Resource ?
                    renderResource((Resource) entry.getValue()) :
                    String.valueOf(entry.getValue());
            if (entry.getKey().equals(AgentsConstants.AgentsConversationEnvironmentKey.ACCEPT_LANGUAGE)){
                locale = I18nUtil.getLocaleWithLanguageTag(value);
            }
            result = result.replace(placeholder, value);
        }

        for (Entry<String, Object> entry : model.entrySet()) {
            String placeholder = "{" + entry.getKey() + "}";
            String value = entry.getValue() instanceof Resource ?
                    renderResource((Resource) entry.getValue()) :
                    String.valueOf(entry.getValue());
            result = result.replace(placeholder, value);
        }

        // To find all placeholders in your strings and replace any missing properties with ""
        Pattern pattern = Pattern.compile(Pattern.quote(startDelimiter) + "(.{1,1000}?)" + Pattern.quote(endDelimiter));
        Matcher matcher = pattern.matcher(result);
        while (matcher.find()) {
            String placeholder = matcher.group();
            if (StringUtils.isNotBlank(I18nUtil.getMessage(matcher.group(1)))){
                result = result.replace(placeholder, I18nUtil.getMessage(matcher.group(1),locale));
            }else if (!dynamicModel.containsKey(matcher.group(1))) {
                result = result.replace(placeholder, "");
            }
        }

        return result;
    }

    private String renderResource(Resource resource) {
        try {
            return resource.getContentAsString(Charset.defaultCharset());
        }
        catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public Message createMessage() {
        return new UserMessage(render());
    }

    @Override
    public Message createMessage(List<Media> mediaList) {
        return new UserMessage(render(), mediaList);
    }

    @Override
    public Message createMessage(Map<String, Object> model) {
        return new UserMessage(render(model));
    }

    @Override
    public Prompt create() {
        return new Prompt(render(new HashMap<>()));
    }

    @Override
    public Prompt create(ChatOptions modelOptions) {
        return new Prompt(render(new HashMap<>()), modelOptions);
    }

    @Override
    public Prompt create(Map<String, Object> model) {
        return new Prompt(render(model));
    }

    @Override
    public Prompt create(Map<String, Object> model, ChatOptions modelOptions) {
        return new Prompt(render(model), modelOptions);
    }

    public Set<String> getInputVariables() {
        Set<String> variables = new HashSet<>();
        int startIndex = 0;
        while (true) {
            int start = template.indexOf(startDelimiter, startIndex);
            if (start == -1) break;

            int end = template.indexOf(endDelimiter, start + startDelimiter.length());
            if (end == -1) break;

            String variable = template.substring(start + startDelimiter.length(), end);
            variables.add(variable.trim());
            startIndex = end + endDelimiter.length();
        }
        return variables;
    }

    private Set<String> getModelKeys(Map<String, Object> model) {
        Set<String> dynamicVariableNames = new HashSet<>(this.dynamicModel.keySet());
        Set<String> modelVariables = new HashSet<>(model.keySet());
        modelVariables.addAll(dynamicVariableNames);
        return modelVariables;
    }

    protected void validate(Map<String, Object> model) {

        Set<String> templateTokens = getInputVariables();
        Set<String> modelKeys = getModelKeys(model);

        // Check if model provides all keys required by the template
        if (!modelKeys.containsAll(templateTokens)) {
            templateTokens.removeAll(modelKeys);
            throw new IllegalStateException(
                    "Not all template variables were replaced. Missing variable names are " + templateTokens);
        }
    }

}
